package com.lmk.dsa.struct.graph;

/**
 * 并查集
 * @author LaoMake
 * @since 1.0
 */
public class UnionFind {

    /** 节点对应的根节点信息 */
    private int[] parent;

    /** 节点所在集合的层数 */
    private int[] rank;

    private int size;

    public UnionFind(int size) {
        this.size = size;
        this.parent = new int[size];
        this.rank = new int[size];
        for (int i = 0; i < size; i++) {
            parent[i] = i;
            rank[i] = 1;
        }
    }

    /**
     * 查找指定节点的根节点
     *
     * @param p
     * @return
     */
    public int find(int p) {
        assert p >= 0 && p < size;
        while (parent[p] != p) {
            parent[p] = parent[parent[p]]; // 路径压缩
            p = parent[p];
        }

        /*// 递归路径压缩
        if(parent[p] != p){
            parent[p] = find(parent[p]);
        }*/

        return p;
    }

    /**
     * 判断两个节点是否连接
     *
     * @param p
     * @param q
     * @return
     */
    public boolean isConnected(int p, int q) {
        return find(p) == find(q);
    }

    /**
     * 合并两个节点
     * @param p
     * @param q
     */
    public void union(int p, int q) {
        int pRoot = find(p);
        int qRoot = find(q);

        if (pRoot == qRoot)
            return;

        // 将元素较少的集合挂载到另外一个集合的根元素下
        if(rank[pRoot] < rank[qRoot]){
            parent[pRoot] = qRoot;
        }else if (rank[pRoot] > rank[qRoot]){
            parent[qRoot] = pRoot;
        }else{
            parent[pRoot] = qRoot;
            rank[qRoot] += 1;
        }
    }
}
